import torch
from torch.utils.data import DataLoader
import torchvision.datasets as dsets
import numpy as np
from tqdm import tqdm

# MNIST dataset
train_dataset = dsets.MNIST(root='./pymnist',  # 选择数据的根目录
                            train=True,  # 选择训练集
                            transform=None,  # 不考虑使用任何数据预处理
                            download=True)  # 从网络上download图片
test_dataset = dsets.MNIST(root='./pymnist',  # 选择数据的根目录
                           train=False,  # 选择测试集
                           transform=None,  # 不考虑使用任何数据预处理
                           download=True)  # 从网络上download图片


def init_netword():
    network = {}
    weight_scale = 1e-3
    network["w1"] = np.random.randn(784, 50) * weight_scale
    network["b1"] = np.ones(50)

    network["w2"] = np.random.randn(50, 100) * weight_scale
    network["b2"] = np.ones(100)

    network["w3"] = np.random.randn(100, 10) * weight_scale
    network['b3'] = np.ones(10)

    return network


def _relu(x):
    return np.maximum(0, x)


def forward(network, x):
    w1, w2, w3 = network['w1'], network['w2'], network['w3']
    b1, b2, b3 = network['b1'], network['b2'], network['b3']
    a1 = x.dot(w1) + b1
    z1 = _relu(a1)
    a2 = z1.dot(w2) + b2
    z2 = _relu(a2)
    a3 = z2.dot(w3) + b3
    y = a3
    return y


if __name__ == '__main__':
    network = init_netword()
    accuaracy_cnt = 0
    x = test_dataset.test_data.numpy().reshape(-1, 28 * 28)
    labels = test_dataset.test_labels.numpy()
    for i in tqdm(range(len(x)), ncols=100):
        y = forward(network, x)
        p = np.argmax(y)
        if p == labels[i]:
            accuaracy_cnt += 1
        # print(i)
    print("Accuracy: " + str(float(accuaracy_cnt) / len(x) * 100) + "%")
